package com.ds.infra.rpc.registry.util;

import com.ds.infra.rpc.registry.constant.ZkConstant;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.utils.CloseableUtils;
import org.apache.zookeeper.CreateMode;

/**
 * Curator工具类，用于创建和管理ZooKeeper客户端
 *
 * @author WeiShaoying
 */
@UtilityClass
@Slf4j
public class CuratorUtils {

    private static volatile CuratorFramework zkClient;

    public static CuratorFramework getZkClient(String zkAddress) {
        if (zkClient == null) {
            synchronized (CuratorUtils.class) {
                if (zkClient == null) {
                    zkClient = CuratorFrameworkFactory.builder()
                            .connectString(zkAddress)
                            .sessionTimeoutMs(ZkConstant.SESSION_TIMEOUT_MS)
                            .connectionTimeoutMs(ZkConstant.CONNECTION_TIMEOUT_MS)
                            .retryPolicy(new ExponentialBackoffRetry(
                                    ZkConstant.RETRY_BASE_SLEEP_MS,
                                    ZkConstant.MAX_RETRIES))
                            .build();
                    zkClient.start();

                    Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                        log.info("Closing ZK client");
                        CloseableUtils.closeQuietly(zkClient);
                    }));
                }
            }
        }
        return zkClient;
    }

    public static void createPersistentNode(CuratorFramework zkClient, String path) throws Exception {
        if (zkClient.checkExists().forPath(path) == null) {
            zkClient.create()
                    .creatingParentsIfNeeded()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(path);
        }
    }
}
